
## R script to make figures for marriage paper from
## outreg2 text files generated by Stata

## NOTE: This code is a slight modification of the main R code
## for the figures. This code was modified to create the main
## figures only to Demography's copyediting style requirements.

## Load required packages:
## readstata13 is a package for reading Stata files
library(readstata13)
## ggplot2 is a popular plotting package by Hadley Wickham
library(ggplot2)
library(gridExtra)

## Write a function to get the figure points from
## the file Stata wrote
get.fig.points <- function(filename, x, group) {
  ## This function takes as arguments
  ## the file name as a character string,
  ## a vector x of the x-values of the figure points,
  ## and the group name as a character string
  results <- read.dta13(paste(
    "/Users/ilundberg/Dropbox/KillewaldLundberg_Demography/Data and code/",
    filename,".dta",sep=""
  ))
  ## Look at just the rows with the coefficient (1),
  ## CI min (5), and CI max (6),
  ## and the columns with graph points
  ## (grepl returns a logical for whether the column name
  ## contains the string "graph")
  results <- results[c(1,5,6), grepl("graph", colnames(results))]
  ## Transpose so that the columns are the coef, CI min, and CI max
  results <- t(results)
  ## Add a first row so we start at 0
  results <- rbind(c(0,0,0),
                   results)
  ## Exponentiate, subtract one, and multiply by 100 for the figures
  results <- 100*(exp(results) - 1)
  ## Rename the columns to clarify what they are
  colnames(results) <- c("pct.change","min","max")
  ## Turn it into a data frame, including the x-values for the plot
  results <- data.frame(x = x, results, group = group)
  return(results)
}

## Function to make the plots
make.plot <- function(results, title, n = NULL, fig.num,
                      min.year = NULL, max.year = NULL,
                      event) {
  ## The ifelse(logical,do1,do2) statement
  ## says to return do1 if logical=T,
  ## otherwise return do2.
  ## If min.year and max.year are manually given, use them
  ## If not, we'll use the min and max of x
  min.year <- ifelse(is.null(min.year), min(results$x), min.year)
  max.year <- ifelse(is.null(max.year), max(results$x), max.year)
  ## Make the figure
  ## The first argument is the data frame with the plot values
  ## The second argument is the aesthetics of the plot,
  ## which tells it what to look for in that data frame
  p <- ggplot(data = results, aes(x = x, y = pct.change, group = group,
                                  color = group, linetype = group)) +
    ## ggplot links commands with +. Here, we tell it
    ## to plot a line plot
    geom_line() +
    ## theme_bw is just an aesthetic option
    theme_bw() +
    ## geom_ribbon adds the CIs. alpha controls the 
    ## transparency of the ribbon
    geom_ribbon(aes(ymin = min, ymax = max, group = group,
                    fill = group),
                alpha = .1, color = NA) +
    ## Set the line types
    scale_linetype_manual(values = c(1,3,5,4)) +
    ## Set the color of the lines
    scale_color_manual(values = c("black","red","blue","green")) +
    ## Set the color of the confidence bands
    scale_fill_manual(values = c("black","red","blue","green")) +
    ## Add a vertical line in the year of marriage
    geom_vline(xintercept = 0,lty = 2) +
    ## Add x and y axis titles
    xlab(paste("\nYears Since",event)) +
    ylab("% Change in Wages\n") +
    ## Change hyphens to en-dashes
    scale_y_continuous(labels = function(x) gsub("-","–",x)) +
    scale_x_continuous(labels = function(x) gsub("-","–",x)) +
    ## Add a plot title.
    ggtitle(if (!is.null(n)) bquote(
      atop(bold(.(title)),
           paste("(",italic(N)," = ",.(prettyNum(n,",")),")"))
    ) else bquote(bold(.(title)))) +
    ## Aesthetic options
    theme(text = element_text(size = 10, family = "Times New Roman"),
          plot.title = element_text(hjust = 0.5),
          axis.title = element_text(face = "bold"),
          legend.position = "none")
  if (length(unique(results$group)) == 1) {
    ggsave(paste("/Users/ilundberg/Dropbox/KillewaldLundberg_Demography/Data and code/CopyeditFigures/",
                 fig.num,".png", sep = ""),
           p, width = 3, height = 3)
  } else {
    leg <- ggplot(data.frame(x = 2.2, y = 0.5 + length(levels(results$group)):1,
                             label = levels(results$group),
                             xmin = 1, xmax = 2,
                             ymin = length(levels(results$group)):1,
                             ymax = 1 + length(levels(results$group)):1),
                  aes(x = x, y = y, label = label,
                      xmin = xmin, xmax = xmax,
                      ymin = ymin, ymax = ymax,
                      fill = label)) +
      geom_text(parse = T, hjust = 0,
                size = 3, family = "Times New Roman") +
      xlim(c(1,8)) + ylim(c(1,4)) +
      geom_rect(aes(fill = label),
                alpha = .2, color = "gray") +
      geom_segment(aes(x = 1.1, xend = 1.9,
                       y = y, yend = y,
                       color = label,
                       linetype = label)) +
      ## Set the line types
      scale_linetype_manual(values = c(1,3,5,4)) +
      ## Set the color of the lines
      scale_color_manual(values = c("black","red","blue","green")) +
      ## Set the color of the confidence bands
      scale_fill_manual(values = c("black","red","blue","green")) +
      theme(legend.position = "none",
            panel.background = element_blank(),
            axis.text = element_blank(),
            axis.ticks = element_blank(),
            axis.title = element_blank())
    arranged <- grid.arrange(p, leg,
                             layout_matrix = rbind(c(1,1),
                                                   c(NA,2)),
                             widths = c(1,9),
                             heights = c(3,1))
    ggsave(paste("/Users/ilundberg/Dropbox/KillewaldLundberg_Demography/Data and code/CopyeditFigures/",
                 fig.num,".png", sep = ""),
           arranged, width = 3, height = 4)
  }
}

#########################
## PRODUCE THE FIGURES ##
#########################

## 1A
## The code here is the style of all the figures. It reads the results
## text file with predicted graph points produced by Stata.
## We need to manually enter the number of respondents in the model (n)
make.plot(results = get.fig.points(filename = "t_a1_c1",
                                   x = c(-5,-1,1,5,10),
                                   group = "Married"),
          title = "a. Marriage (linear spline)",
          n = 4218,
          fig.num = "1a",
          event = "Marriage")

## 1B
make.plot(results = get.fig.points(filename = "t_a2_c1",
                                   x = c(-3,-1,1,5),
                                   group = "Divorced"),
          title = "b. Divorce (linear spline)",
          n = 1755,
          fig.num = "1b",
          event = "Divorce")

## 2A
make.plot(
  results = rbind(
    get.fig.points("t_a3_c1", c(-3,-1,1,5,10),
                   group = paste("Marries~at~age~22~or~under~(italic(N)=='1,545')")),
    get.fig.points("t_a3_c2", c(-3,-1,1,5,10),
                   group = paste0("Marries~at~age~23-26~(italic(N)=='1,301')")),
                     #"Marries at age 23-26 (N = 1,301)"),
    get.fig.points("t_a3_c3", c(-3,-1,1,5,10),
                   group = paste0("Marries~at~age~27+~(italic(N)=='1,372')"))
                     #"Marries at age 27+ (N = 1,372)")
  ),
  title = "a. By Age at Marriage",
  fig.num = "2a",
  event = "Marriage")

## 2B
make.plot(
  results = rbind(
    get.fig.points("t_a3_c4", c(-5,-1,1,5,10),
                   group = "'Non-shotgun'~marriage~(italic(N)=='2,761')"),
    get.fig.points("t_a3_c5", c(-5,-1,1,5,10),
                   group = "Shotgun~marriage~(italic(N)==556)")
  ),
  title = "b. By Shotgun Marriage",
  fig.num = "2b",
  event = "Marriage")

